"""
Reference1: Automatic detection of sub-km craters in high resolution planetary images 
Reference2: Detection of sub-kilometer craters in high resolution planetary images using shape and texture features  
"""

import torch
from torch import nn
import numpy as np
import cv2
from typing import Iterator


class EmamiHoughDetector(nn.Module):
    def __init__(self, *, median_size, th, **kwargs) -> None:
        super().__init__()
        self.median_size = median_size
        self.th = th

    def forward(self, raw: torch.Tensor) -> Iterator[list]:
        """
        API for detecting craters in a pytorch-lightning framework
        Arguments:
            raw: (torch.Tensor), shape=(batch_size, 3, height, width)
        Returns:
            torch.Tensor: shape=(batch_size, n, 4), represent of bounding box in xywh format
        """
        for i, img in enumerate(
            map(lambda item: item.permute(1, 2, 0).cpu().numpy(), raw)
        ):
            img = (img + 1) * 127.5
            img = img.astype("uint8")
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            bboxs = list(self.detect(img))
            yield torch.tensor(bboxs, dtype=torch.float32, device=raw.device).to(
                memory_format=torch.contiguous_format
            )

    def detect(self, image: np.ndarray, ind: int = None) -> Iterator:
        image_copy = cv2.equalizeHist(image)
        image_copy = image.copy()
        # apply median filter on input
        image_copy = cv2.medianBlur(image_copy, self.median_size)
        # Apply the Gaussian filter to the image
        image_copy = cv2.GaussianBlur(image_copy, (5, 5), 0)
        # apply canny edge detector
        image_copy = cv2.Canny(image_copy, 10, 100)

        # cv2.imwrite(f"edge_{ind}.png", image)
        # apply hough transform
        circles = cv2.HoughCircles(image_copy, cv2.HOUGH_GRADIENT, dp=1, minDist=2)
        if circles is not None:
            circles = np.uint16(np.around(circles))
            for circle in circles[0, :]:
                x, y, r = circle
                yield [x - r, y - r, x + r, y + r]
